import os
import glob
import numpy as np
import yaml
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

# 设置防止显存碎片的环境变量（必须放在最前面）
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

def convert_polygon_to_yolo_format(points):
    """
    将四点坐标（左上、左下、右下、右上）转换为YOLO格式的中心点、宽和高
    输入points: [x1,y1,x2,y2,x3,y3,x4,y4]，归一化坐标
    返回: [x_center, y_center, width, height]
    """
    # 提取四个点的坐标
    x1, y1 = points[0], points[1]  # 左上
    x2, y2 = points[2], points[3]  # 左下
    x3, y3 = points[4], points[5]  # 右下
    x4, y4 = points[6], points[7]  # 右上
    
    # 计算中心点
    x_center = (x1 + x2 + x3 + x4) / 4.0
    y_center = (y1 + y2 + y3 + y4) / 4.0
    
    # 计算宽度和高度（使用对角线的平均值近似）
    width = max(np.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2), np.sqrt((x4 - x1) ** 2 + (y4 - y1) ** 2))
    height = max(np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2), np.sqrt((x3 - x4) ** 2 + (y3 - y4) ** 2))
    
    return [x_center, y_center, width, height]

def create_data_yaml(data_dir, output_path="data_armour.yaml"):
    """
    创建数据配置文件
    """
    # 将路径转换为绝对路径
    abs_data_dir = os.path.abspath(data_dir)
    
    # 设置装甲板的类别名称
    names = [
        "BG",   # 0: 蓝色-哨兵
        "B1",   # 1: 蓝色-一号
        "B2",   # 2: 蓝色-二号
        "B3",   # 3: 蓝色-三号
        "B4",   # 4: 蓝色-四号
        "B5",   # 5: 蓝色-五号
        "BO",   # 6: 蓝色-前哨站
        "BB",   # 7: 蓝色-基地
        "RG",   # 8: 红色-哨兵
        "R1",   # 9: 红色-一号
        "R2",   # 10: 红色-二号
        "R3",   # 11: 红色-三号
        "R4",   # 12: 红色-四号
        "R5",   # 13: 红色-五号
        "RO",   # 14: 红色-前哨站
        "RB"    # 15: 红色-基地
    ]
    
    # 获取图片目录
    images_dir = os.path.join(abs_data_dir, "images")
    
    # 使用glob计算图片总数
    total_images = len(glob.glob(os.path.join(images_dir, "*.jpg")))
    
    # 划分训练和验证集（按照9:1的比例）
    train_images = int(total_images * 0.9)
    val_images = total_images - train_images
    
    # 创建数据配置
    data_dict = {
        "path": abs_data_dir,  # 使用绝对路径
        "train": os.path.join(abs_data_dir, "images"),  # 使用绝对路径
        "val": os.path.join(abs_data_dir, "images"),    # 使用绝对路径
        "names": names,
        "nc": len(names)    # 类别数
    }
    
    # 将数据配置写入yaml文件
    with open(output_path, "w") as f:
        yaml.dump(data_dict, f, default_flow_style=False)
    
    print(f"已创建数据配置文件: {output_path}")
    print(f"总图片数: {total_images}，训练集: {train_images}，验证集: {val_images}")
    print(f"数据集路径: {abs_data_dir}")
    print(f"训练数据路径: {data_dict['train']}")
    
    return output_path

def prepare_dataset(data_dir, temp_dir=None):
    """
    准备数据集，直接处理四点标记格式
    """
    # 将路径转换为绝对路径
    data_dir = os.path.abspath(data_dir)
    
    images_dir = os.path.join(data_dir, "images")
    labels_dir = os.path.join(data_dir, "labels")
    
    # 检查目录是否存在
    if not os.path.exists(images_dir) or not os.path.exists(labels_dir):
        raise FileNotFoundError(f"数据目录不存在: {images_dir} 或 {labels_dir}")
    
    image_files = sorted(glob.glob(os.path.join(images_dir, "*.jpg")))
    
    if not image_files:
        raise FileNotFoundError(f"在 {images_dir} 中未找到图像文件")
    
    print(f"找到 {len(image_files)} 张图像文件")
    
    # 创建数据配置文件
    yaml_path = create_data_yaml(data_dir)
    
    return yaml_path

if __name__ == '__main__':
    # 设置数据目录
    data_dir = "./armor_dataset"  # 相对路径
    
    # 获取并显示绝对路径，以便排查问题
    abs_data_dir = os.path.abspath(data_dir)
    print(f"数据集绝对路径: {abs_data_dir}")
    
    # 准备数据集
    data_yaml = prepare_dataset(data_dir)
    
    # 检查本地是否已有模型文件
    model_name = "./yolov12/yolov12x.pt"
    local_model_path = os.path.join(os.getcwd(), model_name)
    
    if os.path.exists(local_model_path):
        print(f"发现本地模型: {local_model_path}")
        model_path = local_model_path
    else:
        # 如果本地没有，则使用URL直接下载
        model_version = "yolov12x"
        model_path = f"https://github.com/sunsmarterjie/yolov12/releases/download/turbo/{model_version}.pt"
        print(f"未找到本地模型，将从网络下载: {model_path}")
    
    # 初始化YOLO模型并加载预训练权重
    print(f"加载模型: {model_path}")
    model = YOLO(model_path, task="detect")
    
    # 训练模型（关键修改部分）
    model.train(
        data=data_yaml,
        imgsz=416,          # 降低分辨率到416
        epochs=100,
        batch=2,            # 减小batch_size到2
        nbs=64,             # 添加nominal batch size
        workers=4,           # 减少数据加载进程
        amp=True,           # 启用混合精度训练
        device='0',
        optimizer='SGD',
        close_mosaic=10,
        resume=False,
        project='runs/train',
        name='armour_detection',
        single_cls=False,
        cache=True,         # 启用数据缓存
        cos_lr=True,
        patience=50,
        save=True,
        save_period=20,
        val=True,
        model=model_path,
        lr0=0.01,          # 明确设置学习率
        momentum=0.937,
        weight_decay=0.0005
    )
    
    # 验证模型
    model.val()